
#-*- coding: utf-8 -*-

#+---------------------------------------------------------------------------+
#|          01001110 01100101 01110100 01111010 01101111 01100010            |
#|                                                                           |
#|               Netzob : Inferring communication protocols                  |
#+---------------------------------------------------------------------------+
#| Copyright (C) 2011-2014 Georges Bossert and Frédéric Guihéry              |
#| This program is free software: you can redistribute it and/or modify      |
#| it under the terms of the GNU General Public License as published by      |
#| the Free Software Foundation, either version 3 of the License, or         |
#| (at your option) any later version.                                       |
#|                                                                           |
#| This program is distributed in the hope that it will be useful,           |
#| but WITHOUT ANY WARRANTY; without even the implied warranty of            |
#| MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the              |
#| GNU General Public License for more details.                              |
#|                                                                           |
#| You should have received a copy of the GNU General Public License         |
#| along with this program. If not, see <http://www.gnu.org/licenses/>.      |
#+---------------------------------------------------------------------------+
#| @url      : http://www.netzob.org                                         |
#| @contact  : contact@netzob.org                                            |
#| @sponsors : Amossys, http://www.amossys.fr                                |
#|             Supélec, http://www.rennes.supelec.fr/ren/rd/cidre/           |
#+---------------------------------------------------------------------------+

#+---------------------------------------------------------------------------+
#| File contributors :                                                       |
#|       - Georges Bossert <georges.bossert (a) supelec.fr>                  |
#|       - Frédéric Guihéry <frederic.guihery (a) amossys.fr>                |
#+---------------------------------------------------------------------------+

#+---------------------------------------------------------------------------+
#| Standard library imports
#+---------------------------------------------------------------------------+
import multiprocessing
import time
import operator

#+---------------------------------------------------------------------------+
#| Local application imports
#+---------------------------------------------------------------------------+
from netzob.Common.Utils.Decorators import typeCheck, NetzobLogger
from netzob.Common.Models.Vocabulary.AbstractField import AbstractField
from netzob.Common.Models.Types.AbstractType import AbstractType
from netzob.Common.Models.Vocabulary.Domain.DomainFactory import DomainFactory
from netzob.Common.Models.Vocabulary.Field import Field
from netzob.Common.Models.Types.TypeConverter import TypeConverter
from netzob.Common.Models.Types.BitArray import BitArray
from netzob.Common.Models.Types.Raw import Raw
from netzob.Common.Models.Types.HexaString import HexaString


@NetzobLogger
class FieldSplitStatic(object):
    """This class allows to update a field structure following
    the static and dynamic portions of data.

    >>> import binascii
    >>> from netzob.all import *
    >>> samples = ["00ff2f00000010",	"00001000000011",	"00fe1f00000012",	"00002000000013", "00ff1f00000014",	"00ff1f00000015",	"00ff2f00000016",	"00fe1f00000017"]
    >>> messages = [RawMessage(data=binascii.unhexlify(sample)) for sample in samples]
    >>> symbol = Symbol(messages=messages)
    >>> symbol.addEncodingFunction(TypeEncodingFunction(HexaString))
    >>> print symbol
    '00ff2f00000010'
    '00001000000011'
    '00fe1f00000012'
    '00002000000013'
    '00ff1f00000014'
    '00ff1f00000015'
    '00ff2f00000016'
    '00fe1f00000017'
    
    >>> fs = FieldSplitStatic()
    >>> fs.execute(symbol)
    >>> print symbol
    '00' | 'ff2f' | '000000' | '10'
    '00' | '0010' | '000000' | '11'
    '00' | 'fe1f' | '000000' | '12'
    '00' | '0020' | '000000' | '13'
    '00' | 'ff1f' | '000000' | '14'
    '00' | 'ff1f' | '000000' | '15'
    '00' | 'ff2f' | '000000' | '16'
    '00' | 'fe1f' | '000000' | '17'

    >>> fs = FieldSplitStatic(mergeAdjacentStaticFields=False, mergeAdjacentDynamicFields=False)
    >>> fs.execute(symbol)
    >>> print symbol
    '00' | 'ff' | '2f' | '00' | '00' | '00' | '10'
    '00' | '00' | '10' | '00' | '00' | '00' | '11'
    '00' | 'fe' | '1f' | '00' | '00' | '00' | '12'
    '00' | '00' | '20' | '00' | '00' | '00' | '13'
    '00' | 'ff' | '1f' | '00' | '00' | '00' | '14'
    '00' | 'ff' | '1f' | '00' | '00' | '00' | '15'
    '00' | 'ff' | '2f' | '00' | '00' | '00' | '16'
    '00' | 'fe' | '1f' | '00' | '00' | '00' | '17'

    >>> fs = FieldSplitStatic(mergeAdjacentStaticFields=True, mergeAdjacentDynamicFields=False)
    >>> fs.execute(symbol)
    >>> print symbol
    '00' | 'ff' | '2f' | '000000' | '10'
    '00' | '00' | '10' | '000000' | '11'
    '00' | 'fe' | '1f' | '000000' | '12'
    '00' | '00' | '20' | '000000' | '13'
    '00' | 'ff' | '1f' | '000000' | '14'
    '00' | 'ff' | '1f' | '000000' | '15'
    '00' | 'ff' | '2f' | '000000' | '16'
    '00' | 'fe' | '1f' | '000000' | '17'

    >>> fs = FieldSplitStatic(mergeAdjacentStaticFields=False, mergeAdjacentDynamicFields=True)
    >>> fs.execute(symbol)
    >>> print symbol
    '00' | 'ff2f' | '00' | '00' | '00' | '10'
    '00' | '0010' | '00' | '00' | '00' | '11'
    '00' | 'fe1f' | '00' | '00' | '00' | '12'
    '00' | '0020' | '00' | '00' | '00' | '13'
    '00' | 'ff1f' | '00' | '00' | '00' | '14'
    '00' | 'ff1f' | '00' | '00' | '00' | '15'
    '00' | 'ff2f' | '00' | '00' | '00' | '16'
    '00' | 'fe1f' | '00' | '00' | '00' | '17'


    We can also plays with the unitsize:
    >>> fs = FieldSplitStatic(AbstractType.UNITSIZE_8, mergeAdjacentDynamicFields=False)
    >>> fs.execute(symbol)
    >>> print symbol
    '00' | 'ff' | '2f' | '000000' | '10'
    '00' | '00' | '10' | '000000' | '11'
    '00' | 'fe' | '1f' | '000000' | '12'
    '00' | '00' | '20' | '000000' | '13'
    '00' | 'ff' | '1f' | '000000' | '14'
    '00' | 'ff' | '1f' | '000000' | '15'
    '00' | 'ff' | '2f' | '000000' | '16'
    '00' | 'fe' | '1f' | '000000' | '17'

    >>> fs = FieldSplitStatic(AbstractType.UNITSIZE_16, mergeAdjacentDynamicFields=False)
    >>> fs.execute(symbol)
    >>> print symbol
    '00ff' | '2f00' | '0000' | '10'
    '0000' | '1000' | '0000' | '11'
    '00fe' | '1f00' | '0000' | '12'
    '0000' | '2000' | '0000' | '13'
    '00ff' | '1f00' | '0000' | '14'
    '00ff' | '1f00' | '0000' | '15'
    '00ff' | '2f00' | '0000' | '16'
    '00fe' | '1f00' | '0000' | '17'

    >>> fs = FieldSplitStatic(AbstractType.UNITSIZE_32, mergeAdjacentDynamicFields=False)
    >>> fs.execute(symbol)
    >>> print symbol
    '00ff2f00' | '000010'
    '00001000' | '000011'
    '00fe1f00' | '000012'
    '00002000' | '000013'
    '00ff1f00' | '000014'
    '00ff1f00' | '000015'
    '00ff2f00' | '000016'
    '00fe1f00' | '000017'

    >>> fs = FieldSplitStatic(AbstractType.UNITSIZE_64, mergeAdjacentDynamicFields=False)
    >>> fs.execute(symbol)
    >>> print symbol
    '00ff2f00000010'
    '00001000000011'
    '00fe1f00000012'
    '00002000000013'
    '00ff1f00000014'
    '00ff1f00000015'
    '00ff2f00000016'
    '00fe1f00000017'

    """

    SUPPORTED_UNITSIZE = [AbstractType.UNITSIZE_8, AbstractType.UNITSIZE_16, AbstractType.UNITSIZE_32, AbstractType.UNITSIZE_64]

    def __init__(self, unitSize=AbstractType.UNITSIZE_8, mergeAdjacentStaticFields=True, mergeAdjacentDynamicFields=True):
        """Constructor.


        :keyword unitSize: the unitsize considered to split
        :type unitSize: :class:`int`
        :keyword mergeAdjacentStaticFields: if set to true, adjacent static fields are merged in a single field
        :type mergeAdjacentStaticFields: :class:`bool`
        :keyword mergeAdjacentDynamicFields: if set to true, adjacent dynamic fields are merged in a single field
        :type mergeAdjacentDynamicFields: :class:`bool`
        """
        self.unitSize = unitSize
        self.mergeAdjacentStaticFields = mergeAdjacentStaticFields
        self.mergeAdjacentDynamicFields = mergeAdjacentDynamicFields

    @typeCheck(AbstractField)
    def execute(self, field):
        """Executes the field edition following the specified messages.
        Children of the specified field will be replaced with new fields.

        :param field: the format definition that will be user
        :type field: :class:`netzob.Common.Models.Vocabulary.AbstractField.AbstractField`
        :raise Exception: if something bad happens
        """

        if field is None:
            raise TypeError("The field cannot be None")
        fieldValues = [TypeConverter.convert(data, Raw, HexaString) for data in field.getValues(encoded=False)]

        if len(fieldValues) == 0:
            raise Exception("No value found in the field.")

        # Retrieve longuest field value
        maxLengthFieldValue = len(max(fieldValues, key=len))

        # definies the step following specified unitsize
        stepUnitsize = self.__computeStepForUnitsize()

        # Vertical identification of variation
        indexedValues = []
        for i in range(0, maxLengthFieldValue, stepUnitsize):
            currentIndexValue = []
            for fieldValue in fieldValues:
                if i < len(fieldValue):
                    currentIndexValue.append(fieldValue[i:min(len(fieldValue), i + stepUnitsize)])
                else:
                    currentIndexValue.append('')
            indexedValues.append(currentIndexValue)

        # If requested, merges the adjacent static fields
        if self.mergeAdjacentStaticFields:
            result = []
            staticSequences = []
            for values in indexedValues:
                if len(set(values)) == 1:
                    # static
                    staticSequences.append(values[0])
                else:
                    # dynamic
                    if len(staticSequences) > 0:
                        result.append([''.join(staticSequences)])
                        staticSequences = []
                    result.append(values)
            if len(staticSequences) > 0:
                result.append([''.join(staticSequences)])
            indexedValues = result

        # If requested, merges the adjacent dynamic fields
        if self.mergeAdjacentDynamicFields:
            result = []
            dynamicSequences = []
            for values in indexedValues:
                if len(set(values)) > 1:
                    # dynamic
                    dynamicSequences.append(values)
                else:
                    # static
                    if len(dynamicSequences) > 0:
                        dynValues = map(None, *dynamicSequences)
                        tmp_result = []
                        for d in dynValues:
                            tmp_result.append(''.join([x if x is not None else '' for x in d]))
                        result.append(tmp_result)
                        dynamicSequences = []
                    result.append(values)
            if len(dynamicSequences) > 0:
                dynValues = map(None, *dynamicSequences)
                tmp_result = []
                for d in dynValues:
                    tmp_result.append(''.join([x if x is not None else '' for x in d]))
                result.append(tmp_result)

            indexedValues = result

        # Create a field for each entry
        newFields = []
        for (i, val) in enumerate(indexedValues):
            fName = "Field-{0}".format(i)            
            fDomain = DomainFactory.normalizeDomain([Raw(TypeConverter.convert(v, HexaString, BitArray)) for v in set(val)])
            newFields.append(Field(domain=fDomain, name=fName))

        # attach encoding functions
        for newField in newFields:
            newField.encodingFunctions = field.encodingFunctions.values()

        field.fields = newFields


    def __computeStepForUnitsize(self):
        """Computes the step following the specified unitsize.

        :return: the step
        :rtype: :class:`int`
        :raise: Exception if unitsize not supported
        """
        if self.unitSize == AbstractType.UNITSIZE_4:
            return 1
        elif self.unitSize == AbstractType.UNITSIZE_8:
            return 2
        elif self.unitSize == AbstractType.UNITSIZE_16:
            return 4
        elif self.unitSize == AbstractType.UNITSIZE_32:
            return 8
        elif self.unitSize == AbstractType.UNITSIZE_64:
            return 16

        else:
            raise Exception("Unitsize not supported, can't compute the step")

    # Static method
    @staticmethod
    def split(field, unitSize=AbstractType.UNITSIZE_8, mergeAdjacentStaticFields=True, mergeAdjacentDynamicFields=True):
        """Split the portion of message in the current field
        following the value variation every unitSize

        :param field : the field to consider when spliting
        :type: :class:`netzob.Common.Models.Vocabulary.AbstractField.AbstractField`
        :keyword mergeAdjacentStaticFields: if set to true, adjacent static fields are merged in a single field
        :type mergeAdjacentStaticFields: :class:`bool`
        :keyword mergeAdjacentDynamicFields: if set to true, adjacent dynamic fields are merged in a single field
        :type mergeAdjacentDynamicFields: :class:`bool`
        :keyword unitSize: the required size of static element to create a static field
        :type unitSize: :class:`int`.
        """

        if field is None:
            raise TypeError("Field cannot be None.")

        if unitSize is None:
            raise TypeError("Unitsize cannot be None.")

        if len(field.messages) < 1:
            raise ValueError("The associated symbol does not contain any message.")

        pSplit = FieldSplitStatic(unitSize, mergeAdjacentStaticFields, mergeAdjacentDynamicFields)
        pSplit.execute(field)

    # Properties

    @property
    def unitSize(self):
        return self.__unitSize

    @unitSize.setter
    @typeCheck(str)
    def unitSize(self, unitSize):
        if unitSize is None:
            raise TypeError("Unitsize cannot be None")

        if unitSize not in FieldSplitStatic.SUPPORTED_UNITSIZE:
            raise ValueError("The specified unitsize is not supported, only {0} are available".format(FieldSplitStatic.SUPPORTED_UNITSIZE))

        self.__unitSize = unitSize

    @property
    def mergeAdjacentStaticFields(self):
        return self.__mergeAdjacentStaticFields

    @mergeAdjacentStaticFields.setter
    @typeCheck(bool)
    def mergeAdjacentStaticFields(self, mergeAdjacentStaticFields):
        if mergeAdjacentStaticFields is None:
            raise TypeError("mergeAdjacentStaticFields cannot be None")

        self.__mergeAdjacentStaticFields = mergeAdjacentStaticFields

    @property
    def mergeAdjacentDynamicFields(self):
        return self.__mergeAdjacentDynamicFields

    @mergeAdjacentDynamicFields.setter
    @typeCheck(bool)
    def mergeAdjacentDynamicFields(self, mergeAdjacentDynamicFields):
        if mergeAdjacentDynamicFields is None:
            raise TypeError("mergeAdjacentDynamicFields cannot be None")

        self.__mergeAdjacentDynamicFields = mergeAdjacentDynamicFields
